package org.example.utils;

import org.elasticsearch.index.query.QueryBuilder;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.script.Script;
import org.elasticsearch.script.ScriptType;

import java.util.Collections;

/**
 * @author ZH
 * @since: 2024/10/15
 */

public class EsQueryUtils {
    /**
     * 构建向量查询
     * 使用 余弦相似度方法cosineSimilarity
     */
    public static QueryBuilder vectorQuery(String field, double[] vector) {
        return QueryBuilders.scriptScoreQuery(
                QueryBuilders.matchAllQuery(),
                new Script(
                        ScriptType.INLINE,
                        "painless",
                        "cosineSimilarity(params.query_vector, doc['" + field + "']) + 1.0",
                        Collections.singletonMap("query_vector", vector)
                )
        );
    }
}
